import copy

import torch

from algorithms.algorithm import Algorithm
from algorithms.large.mocgde import MOCGDE
from problems.problem import Problem


class MOSD(Algorithm):
    parameters = [
        {'label': 'step', 'name': 'step', 'type': 'number', 'description': 'step size', 'step': 0.01,
         'defaultValue': 0.1}
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs,
                 ):
        super().__init__(pop_size, problem, kwargs)
        self.population = problem.init_population(pop_size)
        self.archive = copy.deepcopy(self.population)
        self.step0 = self.step

    def final_output(self):
        return self.archive

    def each_iteration(self):
        for i in range(self.pop_size):
            gk = self.problem.cal_grad(self.population[i].dec)
            gk1 = torch.norm(gk[:, 0])
            gk2 = torch.norm(gk[:, 1])
            if gk1 <= gk2:
                gs = gk[:, 0]
                gl = gk[:, 1]
            else:
                gs = gk[:, 1]
                gl = gk[:, 0]

            if torch.dot(gs, gs - gl) <= 0:
                d = -gs / torch.norm(gs)
            else:
                numerator = (torch.norm(gl) ** 2 - torch.dot(gs, gl))
                denominator = (torch.norm(gs) ** 2 - torch.dot(gs, gl))
                D = ((numerator / denominator) * (-gs)) - gl
                d = D / torch.norm(D)

            offspring_dec = self.population[i].dec + self.step * d
            offspring = self.problem.estimate_population(offspring_dec)[0]
            offspring_obj = offspring.obj
            self.archive = MOCGDE.update_archive(self.archive + offspring, self.pop_size)
            if not torch.any(offspring_obj < self.population[i].obj):
                self.population[i] = self.archive[torch.randint(len(self.archive), (1,)).item()]
                self.step = self.step / 2
            else:
                self.population[i] = offspring
                self.step = min(2 * self.step, self.step0)
